package edu.kit.csl.pisa.models;

/*
This file is part of the PISA Alignment Tool.

Copyright (C) 2013
Karlsruhe Institute of Technology
Cognitive Systems Lab (CSL)
Felix Stahlberg

PISA is free software: you can redistribute it and/or modify
it under the terms of the GNU General Public License as published by
the Free Software Foundation, either version 3 of the License, or
(at your option) any later version.

PISA is distributed in the hope that it will be useful,
but WITHOUT ANY WARRANTY; without even the implied warranty of
MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the
GNU General Public License for more details.

You should have received a copy of the GNU General Public License
along with PISA. If not, see <http://www.gnu.org/licenses/>.
*/

import java.io.BufferedWriter;
import java.io.FileWriter;
import java.io.IOException;
import java.util.Arrays;

import edu.kit.csl.pisa.datatypes.Corpus;
import edu.kit.csl.pisa.datatypes.Dictionary;
import edu.kit.csl.pisa.datatypes.IdentityDictionary;
import edu.kit.csl.pisa.datatypes.SentenceAlignment;
import edu.kit.csl.pisa.datatypes.SentencePair;
import edu.kit.csl.pisa.io.Logger;
import edu.kit.csl.pisa.ui.Configuration;

/**
 * This is the parent class for alignment models like Model 3-phone and 
 * Model 4-phone. All parameters of the models should be encapsulated in the
 * subclasses in order to assure a decoupled training process.
 */
public abstract class AlignmentModel {
	
	/*
	 * The name of this model
	 */
	private String name;
	
	private static final double EPSILON = 1e-50;
	
	/*
	 * Minimum probability as log.
	 */
	private static double logProbSmooth;
	
	/*
	 * Smoothing factor used by normalization.
	 */
	private double smoothFactor;
	
	/**
	 * This arrays stores for each possible fertility f the value log(f!).
	 * This can be used by subclasses when calculating P(a,f|e) instead of
	 * calculating these values.
	 */
	protected static double logFacFert[];
	
	/**
	 * Calculates the term binom(k-phi,phi). This occurs in Model 3P formula.
	 * The term can be computed more efficiently using the logFacFert[] array.
	 * Mathematical background in the documentation. 
	 * 
	 * @param k parameter k
	 * @param phi parameter phi
	 * @return log binom(k-phi, phi), assuming phi <= maxFertility
	 * @throws ArrayIndexOutOfBoundsException if phi > maxFertility
	 */
	protected static double binomSpecial(int k, int phi) {
		double acc = 1;
		final int limit = k - 2*phi;
		if (limit < 0) {
			return logProbSmooth;
		}
		for (int i = k - phi; i > limit; i--) {
			acc *= i;
		}
		return Math.log(acc) - logFacFert[phi];
	}
	
	/**
	 * Calculates P(a,f|e). This method narrows down the whole model to
	 * ideally one mathematical formula. The training process will
	 * maximize the alignments according the probability distribution
	 * of P(a,f|e). Implementers should not use the score given by
	 * {@link SentenceAlignment#getScore()} since it is outdated at this
	 * point.
	 * 
	 * @param a sentence alignment
	 * @return the log probability P(a,f|e)
	 */
	public abstract double calculateAlignmentProbability(SentenceAlignment a);

	/**
	 * This method dumps the alignment model to the file system. For each
	 * model parameter, a separate file should be created. Implementing 
	 * subclasses should use <code>dumpParameter()</code> in order to
	 * maintain a consistent file format. It is also recommended to use
	 * GIZA++ naming scheme and identify files generated by using
	 * dictionaries with the substring <code>.actual.</code>.
	 * 
	 * @param prefix each created file name begins with this string
	 * @param postfix the string is concatenated to all file names
	 * @see #dumpParameter(String, Object[])
	 * @see #dumpParameter(String, Object[], Dictionary[])
	 */
	public abstract void dumpToFilesystem(String prefix, String postfix);
	
	/**
	 * Calculate the perplexity of the model and print it as notice. This 
	 * method should only be called for testing purpose because the time
	 * complexity is linear in the corpus size, not in the number of threads
	 * available (sequential implementation). Note that even without using 
	 * this method, the model perplexity is printed at the end of 
	 * {@link #collectFractionalCounts(Corpus)} without additional 
	 * computational effort. The perplexity is only printed if the logger
	 * output level is at least <code>DEBUG</code>.
	 * 
	 * @param corpus the corpus in respect to which the perplexity should be
	 * 			calculated
	 * @throws NullPointerException if <code>corpus</code> is null
	 * @see edu.kit.csl.pisa.io.Logger#getLevel()
	 * @see edu.kit.csl.pisa.io.Logger#DEBUG
	 */
	public void printPerplexity(Corpus corpus) {
		Logger log = Logger.getSingleton();
		if (log.getLevel() >= Logger.DEBUG) {
			double exp = 0f;
			for (SentencePair pair : corpus) {
				exp += pair.getViterbiScore(); // TODO: Wrong! Use sum!
			}
			Logger.getSingleton().debug(getName() + " e-perplexity is now e^"
					+ (-exp / corpus.size()));
		}
	}
	
	/**
	 * This method is called before {@link #aggregateFractionalCount(
	 * SentenceAlignment, double)} is called the first time. Implementing
	 * subclasses are required to allocate memory in order to be ready to
	 * collect fractional counts, probably in a separate data structure.
	 * 
	 * @see #aggregateFractionalCount(SentenceAlignment, double)
	 * @see #writeBackFractionalCounts()
	 */
	abstract public void initializeFractionalCounts();
	
	/**
	 * This template method (GoF) is the interface for implementing subclasses
	 * to multi-threaded collection of fractional counts. The weight is already
	 * normalized. Therefore, it is equal to P(a|e,f).
	 * 
	 * @param a the alignment of which fractional counts should be collected
	 * @param weight the alignment weight P(a|e,f) of a 
	 * @see #initializeFractionalCounts()
	 * @see #writeBackFractionalCounts()
	 * @throws NullPointerException if no collection is done right now
	 */
	abstract public void aggregateFractionalCount(SentenceAlignment a,
			double weight);
	
	/**
	 * This method is called after a training iteration. The models should
	 * normalize the collected fractional counts and write them back to the
	 * model parameters.
	 * 
	 * @see #initializeFractionalCounts()
	 * @see #aggregateFractionalCount(SentenceAlignment, double)
	 * @throws NullPointerException if no fractional count collection was done
	 */
	abstract public void writeBackFractionalCounts();
	
	   
    /**
     * Transfer model parameters from another model.
     *
     * @param srcModel model from which model parameters should be transferred
     * @throws IllegalArgumentException if the import from the given model
     *                 is not implemented
     */
    public void importModelParameters(AlignmentModel srcModel)
            throws IllegalArgumentException {
        throw new IllegalArgumentException();
    }
	
	/**
	 * Sole constructor.
	 * 
	 * @param name the name of the model. This is used to fetch the right
	 * 			settings for this model.
	 * @see Configuration
	 */
	public AlignmentModel(String name) {
		this.name = name;
		// Initialize logFacFert
		// null fertility could exceed maxFertility, so we have to calculate
		// more values
		logProbSmooth = Math.log(
				Configuration.getSingleton().getFloat("probSmooth"));
		if (!(null == this.getConfigFloat("SmoothFactor"))) {
			smoothFactor = this.getConfigFloat("SmoothFactor");
		}
		Integer maxFert = Configuration.getSingleton().getInteger(
				"maxFertility");
		int logFacFertCount = Math.max(
				Configuration.getSingleton().getInteger("maxTrgtSenLen"),
				maxFert == null ? 0 : maxFert);
		logFacFert = new double[logFacFertCount];
		double acc = logFacFert[0] = logFacFert[1] = 0; // log(1)
		for (int i = 2; i < logFacFertCount; i++) {
			acc += Math.log(i);
			logFacFert[i] = acc;
		}
	}
	
	/**
	 * Get the name of this model.
	 * 
	 * @return the name of this model
	 */
	public String getName() {
		return name;
	}
	
	/**
	 * This wrapper function gets a string configuration parameter from 
	 * Configuration.getSingleton(). The final parameter name is 
	 * composed of the model name and the given key.
	 * 
	 * @param key parameter key (getName() + key have to be a key in
	 * 				Configuration
	 * @see Configuration
	 */
	public String getConfigString(String key) {
		return Configuration.getSingleton().getString(name + key);
	}
	
	/**
	 * This wrapper function gets a integer configuration parameter from 
	 * Configuration.getSingleton(). The final parameter name is 
	 * composed of the model name and the given key.
	 * 
	 * @param key parameter key (getName() + key have to be a key in
	 * 				Configuration
	 * @see Configuration
	 */
	public Integer getConfigInteger(String key) {
		return Configuration.getSingleton().getInteger(name + key);
	}
	
	/**
	 * This wrapper function gets a float configuration parameter from 
	 * Configuration.getSingleton(). The final parameter name is 
	 * composed of the model name and the given key.
	 * 
	 * @param key parameter key (getName() + key have to be a key in
	 * 				Configuration
	 * @see Configuration
	 */
	public Float getConfigFloat(String key) {
		return Configuration.getSingleton().getFloat(name + key);
	}
	
	/**
	 * This method writes a parameter to a file. By parameter, we understand
	 * a multi-dimensional double array. Each combination is written in a
	 * new blank separated line. The last value in each line is the double
	 * value itself.
	 *  
	 * @param fileName name of the file to write
	 * @param parameter the multi-dimensional double array
	 * @throws NullPointerException if one of the arguments are null
	 * @throws ClassCastException if <code>parameter</code> is not a multi-
	 * 		dimensional double array
	 */
	protected void dumpParameter(String fileName, Object[] parameter) {
		dumpParameter(fileName, parameter, new Dictionary[]{
			null, null, null, null, null, null, null, null, null
		});
	}
	
	/**
	 * This method writes a parameter to a file. By parameter, we understand
	 * a multi-dimensional double array. Each combination is written in a
	 * new blank separated line. The last value in each line is the double
	 * value itself. Before the line is written, some elements may be replaced
	 * by their string representations. For defining the replacements, give
	 * a appropriate dictionary to look up for the dimension you wish. Note
	 * that files generated with this method are similar to those generated
	 * by GIZA++.
	 *  
	 * @param fileName name of the file to write
	 * @param parameter the multi-dimensional double array
	 * @param dictionaries dictionaries used to replace values in certain
	 * 		columns (give null for those dimensions you do not need one)
	 * @throws NullPointerException if one of the arguments are null
	 * @throws ArrayIndexOutOfBoundsException if <code>dictionaries</code> has
	 * 				not enough elements
	 * @throws ClassCastException if <code>parameter</code> is not a multi-
	 * 		dimensional double array
	 */
	protected void dumpParameter(String fileName, Object[] parameter,
			Dictionary[] dictionaries) {
		for (int i = 0; i < dictionaries.length; i++) {
			if (dictionaries[i] == null) {
				dictionaries[i] = new IdentityDictionary();
			}
		}
		try {
			FileWriter fs = new FileWriter(fileName);
			BufferedWriter out = new BufferedWriter(fs);
			writeParameter(out, "", parameter, dictionaries, 0);
			out.close();
		} catch (IOException e) {
			Logger.getSingleton().error("Error while dumping model to "
					+ fileName + ": " + e.getMessage());
		}
	}
	
	/* (non-Javadoc)
	 * Recursive method for writing multi-dimensional double arrays to
	 * a buffered writer.
	 */
	private void writeParameter(BufferedWriter out, String prefix,
			Object[] param, Dictionary[] dicts, int dictPos) 
					throws IOException {
		int nextDictPos = dictPos + 1;
		Dictionary dict = dicts[dictPos];
		Dictionary nextDict = dicts[nextDictPos];
		for (int i = 0; i < param.length; i++) {
			if (param[i] instanceof Object[]) {
				// Recursive
				writeParameter(out, prefix + dict.get(i) + " ",
						(Object[]) param[i], dicts, nextDictPos);
			} else {
				double[] arr = (double[]) param[i];
				for (int j = 0; j < arr.length; j++) {
					out.write(prefix +  dict.get(i) + " "
							+ nextDict.get(j) + " " + Math.exp(arr[j]) + "\n");
				}
			}
		}
	}
	
	/**
	 * Helper method usually used after collecting fractional counts. Takes
	 * an array and normalize its values by the sum of all elements in the
	 * array. Furthermore, it takes the logarithm. This will convert the 
	 * fractional counts into a probability distribution (stored as log).
	 * 
	 * @param arr array which should be normalized
	 */
	public void normalize(double[] a) {
		double acc = sum(a);
		if (acc < EPSILON) {
			Arrays.fill(a, -Math.log(a.length));
			//Arrays.fill(a, Math.log(0.99));
			return;
		}
		final double add = smoothFactor / a.length;
		final double fac = (1.0 - smoothFactor) / acc;
		//final double logAcc = Math.log(acc);
		for (int i = 0; i < a.length; i++) {
			a[i] = Math.max(logProbSmooth, Math.log(add + fac * a[i]));
			//a[i] = Math.max(logProbSmooth, Math.log(a[i] / acc));
		}
	}
	
	/**
	 * Calculates \sum(a_i) with a recursive D&C algorithm to avoid 
	 * underflows. The a_i are stored in the argument a.
	 */
	public static double sum(double[] a) {
		if (a.length == 0) {
			return 0.0;
		}
		return sum(a, 0, a.length);
	}
	
	/*
	 * Helper function for logSum(double[]). Subsum from start (inclusive) to
	 * end (exclusive).
	 */
	protected static double sum(double[] a, int start, int end) {
		if (end - start == 1) {
			return a[start];
		}
		int med = start + (end - start) / 2;
		return sum(a, start, med) + sum(a, med, end);
	}
	
	/**
	 * Calculates \sum(a_i) with a recursive D&C algorithm to avoid 
	 * underflows. The a_i are stored in the argument a.
	 */
	public static double sum(Double[] a) {
		if (a.length == 0) {
			return 0.0;
		}
		return sum(a, 0, a.length);
	}
	
	/*
	 * Helper function for logSum(double[]). Subsum from start (inclusive) to
	 * end (exclusive).
	 */
	protected static double sum(Double[] a, int start, int end) {
		if (end - start == 1) {
			return a[start];
		}
		int med = start + (end - start) / 2;
		return sum(a, start, med) + sum(a, med, end);
	}
	
	/**
	 * Models often output the size of their parameter tables. This method
	 * converts the raw number of doubles to a human readable string. Code is
	 * mainly taken from
	 * {@link http://stackoverflow.com/questions/3758606/how-to-convert-byte-size-into-human-readable-format-in-java}.
	 * 
	 * @param doubles number of double variables to allocate
	 * @return human readable string describing the size
	 */
	protected static String humanReadableDoubleCount(long doubles) {
		long bytes = doubles * 8;
	    final int unit = 1024;
	    if (bytes < unit) return bytes + " B";
	    int exp = (int) (Math.log(bytes) / Math.log(unit));
	    String pre = "KMGTPE".charAt(exp-1) + "i";
	    return String.format("%.1f %sB", bytes / Math.pow(unit, exp), pre);
	}
}
